#include <string.h>
#include <openssl/obj_mac.h>
#include <openssl/sm3.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/err.h>
#include "crypto.h"
#include "utils.h"

/**
 * @brief  Generate a pair of SM2 key
 *
 * @param  [in]  priv_key_file    The private key file name
 * @param  [in]  pubkey           The address of the public key memory
 * @param  [in]  priv_key         The address of the private key memory
 *
 *
 * @returns EC_KEY structure for success, NULL for failure
 */
EC_KEY* get_sm2_key_pair(char * priv_key_file, sm2_pubkey* pubkey, unsigned char* priv_key)
{
    EC_KEY*         ec_key          = NULL;
    const BIGNUM*   priv_key_bignum = NULL;
    EC_POINT* pubkey_ec_point       = NULL;
    const EC_GROUP*       ecc_group = NULL;
    BIGNUM*         x               = NULL;
    BIGNUM*         y               = NULL;
    int             privkey_len     = 0;
    int             pubx_len        = 0;
    int             puby_len        = 0;
    FILE            *fp             = NULL;

    if (priv_key == NULL || pubkey == NULL) {
        printf("invalid pubkey or priv_key\n");
        return NULL;
    }

    fp = fopen(priv_key_file, "rb");
    if (fp == NULL) {
        printf("Failed to open private key file %s\n", priv_key_file);
        return NULL;
    }

    ec_key = PEM_read_ECPrivateKey(fp, NULL, NULL, NULL);
    fclose(fp);

    if (ec_key == NULL) {
        printf("fail to read key from file %s\n", priv_key_file);
        return NULL;
    }


    x        = BN_new();
    y        = BN_new();

    ecc_group = EC_KEY_get0_group(ec_key);
    if (ecc_group == NULL) {
        fprintf(stderr, "Failed to get EC group.\n");
        goto finish;
    }
    memset(priv_key, 0, ECC_KEY_SIZE);
    priv_key_bignum = EC_KEY_get0_private_key(ec_key);
    privkey_len     = BN_bn2bin(priv_key_bignum, priv_key);

    pubkey_ec_point = EC_POINT_new(ecc_group);
    if (pubkey_ec_point == NULL) {
        fprintf(stderr, "Failed to create EC point for public key.\n");
        ec_key = NULL;
        goto finish;
    }
    if (EC_POINT_mul(ecc_group, pubkey_ec_point, priv_key_bignum, NULL, NULL, NULL) != 1) {
        printf("get_sm2_key_pair: fail to calculate the public key\n");
        ec_key = NULL;
        goto finish;
    }


    EC_POINT_get_affine_coordinates_GFp(ecc_group, pubkey_ec_point, x, y, NULL);
    memset(pubkey->qx, 0, ECC_KEY_SIZE);
    memset(pubkey->qy, 0, ECC_KEY_SIZE);
    pubx_len      = BN_bn2bin(x, pubkey->qx);
    puby_len      = BN_bn2bin(y, pubkey->qy);
    pubkey->curve = CURVE_ID_TYPE_SM2_256;


    if ((ECC_KEY_SIZE != privkey_len) || (ECC_KEY_SIZE != pubx_len) || (ECC_KEY_SIZE != puby_len)) {
        printf("get_sm2_key_pair: fail to get the private key or public key\n");
        ec_key = NULL;
    }

finish:
    if (x)
        BN_free(x);

    if (y)
        BN_free(y);

    return ec_key;
}

/*
 * KDF(Z, key, klen)
 * Z: the input bit string
 * key: the output key bit string
 * klen: bit number of the key
 *
 * Split klen bits into blocks of 32 bytes.
 *
 * n = ALIGN_UP(klen/256)
 * ct is a 32 bit string starting from 0x00000001
 *
 * for i from 1 to n
 *   Hi = SM3(Z||ct)
 *   ct ++
 *
 * if klen is multiple of 256
 *  key = H1||H2||...||Hn
 * else
 *  key = H1||H2||...|| [Hn's left klen - (n-1)*256 bit]
 *
 */
#define KDF_SM3_OUTPUT_BLOCK_BYTES          32
#define KDF_SM3_COUNTER_STRING_BYTES        8
#define KDF_SM3_COUNTER_BIN_BYTES           4

/**
 * @brief  change a string of Ascii(0--f) to BYTE str
 *
 * Eg:"1A2B3C4D" (length of 8) will be trasform to byte string 0x1A2B3C4D  (length will be 4)
 * @param  [in] src: source string
 * @param  [in] lsrc: length of source string
 * @param  [out]out: output BYTE str
 * @param  [out]lout: length of output BYTE str
 *
 * @return
 * - success: 0
 * - failure: other errcode
 */
int hex_str_2_u8_bin(char* src, unsigned long lsrc, unsigned char* out, unsigned long* lout)
{
    unsigned int i, j = 0;

    if (src == NULL || out == NULL) {
        return -1;
    }


    if (lsrc != KDF_SM3_COUNTER_STRING_BYTES) {
        printf("hex_str_2_u8_bin invalid input length\n");
        return -1;
    }

    if ((lsrc % 2) != 0) {
        return -1;
    }

    for (i = 0; i < lsrc; i ++) {
        int tmp            = 0;
        unsigned int high_half_byte = 0;
        unsigned int low_half_byte = 0;
        unsigned int byte = 0;


        if (src[i] >= 0x30 && src[i] <= 0x39) {
            /* '0' - '9' */
            byte = src[i] - 0x30;
        } else if (src[i] >= 0x41 && src[i] <= 0x46) {
            /* 'A' - 'F' */
            byte = src[i] - 0x37;
        } else if (src[i] >= 0x61 && src[i] <= 0x66) {
            /* 'a' - 'f' */
            byte = src[i] - 0x57;
        } else {
            return -1;
        }

        if (i%2 == 0) {
            high_half_byte = byte;
        } else {
            low_half_byte  = byte;
            tmp = (high_half_byte << 4) | low_half_byte;
            out[j] = tmp;
            j ++;
        }
    }

    if (lout)
        *lout = j;

    return 0;
}

/*
 *
 * @brief calculate the KDF block where the output length <= 32 bytes
 *
 * @param  [in/out]  out          The output buffer address
 * @param  [in]      outlen       The output buffer length (<= 32 bytes)
 * @param  [in]      in           The input buffer address
 * @param  [in]      inlen        The input buffer length (in bytes)
 * @param  [in]      count_index  The count
 *
 * @returns 0 for success, -1 failure
 *
 */

static int sm3_kdf_block(unsigned char* out, unsigned long outlen,
                  unsigned char* in, unsigned long inlen,
                  int count_index)
{
    char count_str[KDF_SM3_COUNTER_STRING_BYTES+1]    = {0};
    uint8_t count_bin[KDF_SM3_COUNTER_BIN_BYTES]    = {0};
    uint8_t  *p                     = NULL;
    uint8_t   *new_buffer           = NULL;
    unsigned long   len_ct_bin      = 0;
    uint8_t tmp_buff[KDF_SM3_OUTPUT_BLOCK_BYTES] = {0};
    int   ret             = -1;

    if (out == NULL || in == NULL) {
    }
    if (outlen > KDF_SM3_OUTPUT_BLOCK_BYTES) {
        printf("sm3_kdf_block invalid outlen %lu\n", outlen);
        return -1;
    }

    sprintf(count_str, "%08x", (unsigned int)count_index);
    ret    = hex_str_2_u8_bin(count_str, strlen(count_str), count_bin, &len_ct_bin);
    if (ret) {
        return ret;
    }

    if (len_ct_bin != KDF_SM3_COUNTER_BIN_BYTES) {
        printf(" invalid count bin len %lu\n", len_ct_bin);
        return ret;
    }

    //allocatte a new buffer including input buffer + count string
    new_buffer = calloc(1, inlen + len_ct_bin);
    if (!new_buffer) {
        printf("fail to allocate %lu bytes\n", inlen + len_ct_bin);
        return -1;
    }
    p = new_buffer;
    memcpy(p, in, inlen);
    p += inlen;
    memcpy(p, count_bin, KDF_SM3_COUNTER_BIN_BYTES);
    sm3(new_buffer, inlen + KDF_SM3_COUNTER_BIN_BYTES, tmp_buff);

    memcpy(out, tmp_buff, outlen);

    if (new_buffer)
        free(new_buffer);

    return 0;
}

/*
 *
 * @brief calculate the KDF SM3 output
 *
 * Divide the output to blocks and calculate one by one
 *
 * @param  [in/out]  out          The output buffer address
 * @param  [in]      outlen       The output length
 * @param  [in]      in           The input buffer address
 * @param  [in]      inlen        The input buffer length (in bytes)
 *
 * @returns 0 for success, -1 failure
 *
 */
int sm3_kdf(unsigned char* out, unsigned long outlen, unsigned char* in, unsigned long inlen)
{
    int ret = 0;
    unsigned long klen = 0;
    unsigned long ct             = 0;
    unsigned long max_iter       = 0;
    unsigned long reminder       = 0;

    klen = outlen;

    if (out == NULL || in == NULL || inlen == 0 ||klen == 0) {
        return -1;
    }

    max_iter       = (klen) / KDF_SM3_OUTPUT_BLOCK_BYTES;


    /* reminder bytes */
    reminder = klen - max_iter*KDF_SM3_OUTPUT_BLOCK_BYTES;

    for (ct = 0; ct < max_iter; ct++) {
        ret = sm3_kdf_block(out + ct*KDF_SM3_OUTPUT_BLOCK_BYTES, KDF_SM3_OUTPUT_BLOCK_BYTES,
                          in, inlen, ct + 1);
    }

    if (ret) {
        return ret;
    }

    /* reminder */
    if (reminder)
        ret = sm3_kdf_block(out + ct*KDF_SM3_OUTPUT_BLOCK_BYTES, reminder,
                            in, inlen, ct + 1);

    return ret;
}

int kdf_gen_key(unsigned char* key, unsigned int key_size,
                unsigned char* label, unsigned int label_size,
                unsigned char* context, unsigned int context_size,
                unsigned char* result, unsigned int result_size)
{
    int            internal_ret = 0;
    uint64_t       z_len = 0;
    unsigned char *z = NULL;
    unsigned char  cat_symbol = 0x00;
    unsigned char  sha_key[SHA_256_SIZE] = {0};

    if (!key || !label || !result) {
        printf("input invalid parameter.\n");
        return -1;
    }

    z_len = key_size + label_size + sizeof(cat_symbol) + context_size + SHA_256_SIZE;

    z     = (unsigned char*)malloc(z_len);
    if (NULL == z) {
        printf("raw msg malloc z failed.\n");
        return -1;
    }

    memset(z, 0, z_len);
    memcpy(z, key, key_size);
    memcpy(z + key_size, label, label_size);
    memcpy(z + key_size + label_size, &cat_symbol, sizeof(cat_symbol));

    if (NULL != context) {
        memcpy(z + key_size + label_size + sizeof(cat_symbol), context, context_size);
    }

    internal_ret =  sm3_kdf(sha_key, SHA_256_SIZE, z, z_len - SHA_256_SIZE);

    if (internal_ret != 0) {
        printf("kdf_gen_key step 1 sm3_kdf \n");
        return -1;
    }

    memcpy(z + z_len - SHA_256_SIZE, sha_key, SHA_256_SIZE);
    internal_ret = sm3_kdf(result, result_size, z, z_len);
    if (internal_ret != 0) {
        printf("kdf_gen_key step 2 sm3_kdf \n");
        return -1;
    }

    if (z != NULL) {
        free(z);
    }

    return 0;
}

/*
 *
 * @brief do the SM4 CTR mode encryption
 *
 * @param  [in]  enckey       The key for operation
 * @param  [in]  enciv        The IV for operation
 * @param  [in]  input_buf    The data buffer to be encrypted
 * @param  [in]  input_len    The input buffer length
 * @param  [out]  output_buf  The output data
 *
 * @returns 0 for success, -1 failure
 *
 */
int sm4_ctr128_encrypt(const unsigned char *enckey, const unsigned char *enciv,
                    const unsigned char *input_buf, int input_len,
                    unsigned char *output_buf)
{
    EVP_CIPHER_CTX *ctx = NULL;
    int len             = 0;
    int ciphertext_len  = 0;
    int ret             = -1;

    if (!enckey || !enciv || !input_buf || input_len == 0 ||!output_buf) {
        printf("sm4_ctr128_encrypt:invalid parameter\n");
        goto finish;
    }

    if (!(ctx = EVP_CIPHER_CTX_new())) {
        goto finish;
    }

    if (1 != EVP_EncryptInit_ex(ctx, EVP_sms4_ctr(), NULL, enckey, enciv)) {
        printf("sm4_ctr128_encrypt:EVP_EncryptInit_ex error\n");
        goto finish;
    }

    if (1 != EVP_EncryptUpdate(ctx, output_buf, &len, input_buf, input_len)) {
        printf("sm4_ctr128_encrypt:EVP_EncrypUpdate error\n");
        goto finish;
    }

    ciphertext_len = len;

    if (1 != EVP_EncryptFinal_ex(ctx, output_buf + len, &len)) {
        printf("sm4_ctr128_encrypt:EVP_EncrypFinal_ex error\n");
        goto finish;
    }
    ciphertext_len += len;


    if (ciphertext_len != input_len)
        ret = -1;
    else
        ret = 0;

finish:
    EVP_CIPHER_CTX_free(ctx);

    return ret;
}

/*
 *
 * @brief generate random values
 *
 * @param  [in]  bytes           The memory address
 * @param  [in]  num_bytes       The memory length
 *
 *
 * TODO: replace rand() by 'rdrand' instruction
 *
 */
void gen_random_bytes(void* bytes, size_t num_bytes)
{
    uint8_t* addr = (uint8_t*)bytes;

    srand((unsigned int)time(NULL));

    while (num_bytes--) {
        *addr = (uint8_t)(rand() & 0xff);
         addr ++;
    }
}

/*
 * @brief convert the SM2 public key coordinate to EC_KEY structure
 *
 * @param [in]  coord    The X,Y coordinate
 * @param [out] pkey     The pointer to the EC_KEY pointer
 *
 * @returns 0 for success, -1 failure
 */

int sm2_pubkey_coords_to_ec_key(ecc_point * coord, EC_KEY **pkey)
{
    EC_GROUP *group = NULL;
    EC_KEY *key     = NULL;
    EC_POINT *point = NULL;
    BIGNUM *bn_x    = NULL;
    BIGNUM *bn_y    = NULL;
    int  ret        = -1;

    OpenSSL_add_all_algorithms();

    if (!coord || !pkey) {
        printf("sm2_pubkey_coords_to_ec_key NULL coord or NULL key \n");
        goto finish;
    }

    group = EC_GROUP_new_by_curve_name(NID_sm2p256v1);
    if (!group) {
        printf("sm2_pubkey_coords_to_ec_key fail to allocate group\n");
        goto finish;
    }

    key = EC_KEY_new();
    if (!key) {
        printf("sm2_pubkey_coords_to_ec_key fail to allocate key\n");
        goto finish;
    }

    if (EC_KEY_set_group(key, group) != 1) {
        printf("sm2_pubkey_coords_to_ec_key fail to allocate key\n");
        goto finish;
    }

    point = EC_POINT_new(group);
    if (!point) {
        printf("sm2_pubkey_coords_to_ec_key fail to allocate point\n");
        goto finish;
    }

    bn_x = BN_new();
    bn_y = BN_new();
    if (!BN_bin2bn(coord->qx, ECC_KEY_SIZE, bn_x) || !BN_bin2bn(coord->qy, ECC_KEY_SIZE, bn_y)) {
        printf("sm2_pubkey_coords_to_ec_key BN_bin2bn error \n");
        goto finish;

    }

    if (EC_POINT_set_affine_coordinates_GFp(group, point, bn_x, bn_y, NULL) != 1) {
        printf("sm2_pubkey_coords_to_ec_key EC_POINT_set_affine_coordinates_GFp error \n");
        goto finish;
    }

    if (EC_KEY_set_public_key(key, point) != 1) {
        printf("sm2_pubkey_coords_to_ec_key EC_KEY_set_public_key error \n");
        goto finish;
    }

    *pkey = key;

     ret = 0;

finish:
    if (point)
        EC_POINT_free(point);

    if (bn_x)
        BN_free(bn_x);

    if (bn_y)
        BN_free(bn_y);

    if (group)
        EC_GROUP_free(group);

    if ((ret != 0) && key)
        EC_KEY_free(key);


    return ret;
}


/*
 * @brief read the public key coordinates from cert
 *
 * @param [in]  cert   The csv cert structure
 * @param [out] point  The XY coordinates
 *
 * @returns 0 for success, -1 failure
 */
int get_public_key_in_cert(csv_cert *cert, ecc_point* point)
{
    if (!cert || !point) {
        return -1;
    }


    point->curve = cert->pub_key.curve;
    memcpy(point->qx, cert->pub_key.qx, sizeof(point->qx));
    memcpy(point->qy, cert->pub_key.qy, sizeof(point->qy));

    /* little endian in cert, big endian in ecc_point  */
    reverse_bytes(point->qx, ECC_KEY_SIZE);
    reverse_bytes(point->qy, ECC_KEY_SIZE);

    return 0;
}


/*
 * @brief write the public key coordinates to cert
 *
 * @param [in]  cert   The csv cert structure
 * @param [in]  pubkey The public key including XY coordinates
 *
 * @returns 0 for success, -1 failure
 */
int fill_pubkey_in_cert(csv_cert *cert, sm2_pubkey *pubkey)
{
    if (!cert || !pubkey) {
        return -1;
    }

    memcpy(&cert->pub_key, pubkey, sizeof(*pubkey));

    /* big endian in pubkey, little endian in cert */
    reverse_bytes(cert->pub_key.qx, ECC_KEY_SIZE);
    reverse_bytes(cert->pub_key.qy, ECC_KEY_SIZE);

    return 0;
}
